﻿// <copyright file="TTableStorageDataContext.cs" company="N/A">
//     Copyright (c) Nicholas Barrett and Benjamin Kasten.  All rights reserved.
//
//     This program is free software; you can redistribute it and/or modify it
//     under the terms of the GNU General Public License as published by the
//     Free Software Foundation; either version 2 of the License, or (at your
//     option) any later version.
//
//     This program is distributed in the hope that it will be useful, but
//     WITHOUT ANY WARRANTY; without even the implied warranty of
//     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General
//     Public License for more details.
//
//     You should have received a copy of the GNU General Public License along
//     with this program; if not, write to the Free Software Foundation, 
//     Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
// </copyright>

namespace BetterAzureTableStorage
{
    using System;
    using System.Collections.Generic;
    using System.Data.Services.Client;
    using System.Linq;
    using System.Net;
    using System.Reflection;
    using System.Runtime.Serialization;
    using System.Threading;
    using Microsoft.WindowsAzure;
    using Microsoft.WindowsAzure.StorageClient;
   
    /// <summary>
    /// The TableStorageDataContext represents the runtime context of the table service. 
    /// </summary>
    /// <remarks>
    /// This class is similar to DataServiceContext, but tailored for Table Storage.
    /// It handles tracking/reading/writing to/from table storage.
    /// </remarks>
    public class TTableStorageDataContext
    {
        /// <summary>
        /// The number of potentially billable transactions performed by this context.
        /// </summary>
        private ulong fTransactionCount;

        /// <summary>
        /// Gets the transaction count.
        /// </summary>
        public ulong TransactionCount
        {
            get
            {
                return this.fTransactionCount;
            }
        }

        /// <summary>
        /// Gets or sets the retry policy.
        /// </summary>
        /// <value>The retry policy.</value>
        public RetryPolicy RetryPolicy { get; set; }

        /// <summary>
        /// Delegate type for DoGetTableName
        /// </summary>
        /// <param name="entityType">The type of the entity to get table information about.</param>
        /// <returns>A delegate used to get the table information</returns>
        public delegate TTableStorageInfo TDoGetTableInfo(Type entityType);

        /// <summary>
        /// Gets or sets the delegate to lookup the table name of a type
        /// </summary>
        public TDoGetTableInfo DoGetTableInfo { get; set; }

        /// <summary>
        /// Gets or sets an action to be called when incrementing the transaction count.
        /// </summary>
        public Action IncrementingTransactionCount { get; set; }

        /// <summary>
        /// Gets or sets the service root.
        /// </summary>
        /// <value>The service root.</value>
        private Uri ServiceRoot { get; set; }

        /// <summary>
        /// Gets or sets the credentials used to sign requests to the table service.
        /// </summary>
        private StorageCredentials Credentials { get; set; }

        /// <summary>
        /// Gets or sets the removed entities.
        /// </summary>
        /// <remarks>Entities added to this list will be removed from Entities after processing table requests.</remarks>
        private List<TTableEntityTransactionInfo> RemovedEntities { get; set; }

        /// <summary>
        /// Gets or sets the dictionary of entities and the state they are in.
        /// </summary>
        private List<TTableEntityTransactionInfo> Entities { get; set; }

        /// <summary>
        /// The lock object for the fTypeTableInfo dictionary and fTableNameTableInfo dictionary
        /// </summary>
        private static object fTableInfoLockObject = new object();

        /// <summary>
        /// The lock object for the FCreateTables dictionary
        /// </summary>
        private static object fCreatedTablesLockObject = new object();

        /// <summary>
        /// Static dictionary for created table names by service root
        /// </summary>
        private static Dictionary<Uri, List<string>> fCreatedTables = new Dictionary<Uri, List<string>>();

        /// <summary>
        /// Static dictionary for table storage info by type
        /// </summary>
        private static Dictionary<Type, TTableStorageInfo> fTypeTableInfo = new Dictionary<Type, TTableStorageInfo>();

        /// <summary>
        /// Static dictionary for table storage info by table name
        /// </summary>
        private static Dictionary<string, List<TTableStorageInfo>> fTableNameTableInfo = new Dictionary<string, List<TTableStorageInfo>>();

        /// <summary>
        /// Initializes a new instance of the <see cref="TTableStorageDataContext"/> class.
        /// </summary>
        /// <param name="serviceRoot">The base Table Service URI.</param>
        /// <param name="credentials">The credentials.</param>
        /// <param name="getTableInfoDel">The delegate that is used when getting the table information.</param>
        public TTableStorageDataContext(string serviceRoot, StorageCredentials credentials, TDoGetTableInfo getTableInfoDel)
            : this(new Uri(serviceRoot), credentials, getTableInfoDel)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="TTableStorageDataContext"/> class.
        /// </summary>
        /// <param name="serviceRoot">The base Table Service URI.</param>
        /// <param name="credentials">The credentials.</param>
        /// <param name="getTableInfoDel">The delegate that is used when getting the table information.</param>
        public TTableStorageDataContext(Uri serviceRoot, StorageCredentials credentials, TDoGetTableInfo getTableInfoDel)
        {
            this.ServiceRoot = serviceRoot;
            this.Credentials = credentials;
            this.DoGetTableInfo = getTableInfoDel;

            this.RetryPolicy = RetryPolicies.Retry(3, new TimeSpan(0, 0, 30));

            this.RemovedEntities = new List<TTableEntityTransactionInfo>();
            this.Entities = new List<TTableEntityTransactionInfo>();

            if (!fCreatedTables.ContainsKey(serviceRoot))
            {
                CloudStorageAccount account = new CloudStorageAccount(
                    this.Credentials,
                    null,
                    null,
                    this.ServiceRoot);

                CloudTableClient client = account.CreateCloudTableClient();
                List<string> createdTables = new List<string>(client.ListTables());

                lock (fCreatedTablesLockObject)
                {
                    if (!fCreatedTables.ContainsKey(serviceRoot))
                    {
                        fCreatedTables.Add(this.ServiceRoot, createdTables);
                    }
                }
            }
        }

        /// <summary>
        /// Creates a table using a TTableStorageEngine if the table is not found in Azure Table Storage.
        /// </summary>
        /// <param name="tableName">The name of the table being checked and created if necessary.</param>
        public void CreateTableIfNotExists(string tableName)
        {
            List<string> tableNames;
            if (fCreatedTables.TryGetValue(this.ServiceRoot, out tableNames))
            {
                if (!tableNames.Contains(tableName))
                {
                    lock (fCreatedTablesLockObject)
                    {
                        if (!tableNames.Contains(tableName))
                        {
                            TTableStorageEngine engine = new TTableStorageEngine(this.ServiceRoot, this.Credentials, null, tableName, this.RetryPolicy, this.IncrementTransactionCount);
                            if (engine.CreateTable())
                            {
                                fCreatedTables[this.ServiceRoot].Add(tableName);
                            }
                        }
                    }
                }
            }
        }

        /// <summary>
        /// Determines whether the state of the specified entity is known by the context.
        /// </summary>
        /// <param name="entity">The entity.</param>
        /// <returns>
        ///     <c>true</c> if the state of the specified entity is known by the context; otherwise, <c>false</c>.
        /// </returns>
        public bool IsTrackingEntity(object entity)
        {
            return this.Entities.Count(info => info.Entity == entity) > 0;
        }

        /// <summary>
        /// Marks the entity as Modified
        /// </summary>
        /// <param name="entity">The entity to mark as modified.</param>
        public void UpdateObject(ITableEntity entity)
        {
            this.AddOrChangeState(entity, EntityStates.Modified);
        }

        /// <summary>
        /// Adds the entity.
        /// </summary>
        /// <param name="entity">The entity to add.</param>
        public void AddObject(ITableEntity entity)
        {
            this.AddOrChangeState(entity, EntityStates.Added).SetName = this.GetTableInfo(entity.GetType()).TableName;
        }

        /// <summary>
        /// Mark an entity for deletion.
        /// </summary>
        /// <param name="entity">The entity.</param>
        public void DeleteObject(ITableEntity entity)
        {
            this.AddOrChangeState(entity, EntityStates.Deleted);
        }

        /// <summary>
        /// Saves the changes.
        /// </summary>
        public void SaveChanges()
        {
            this.SaveChanges(SaveChangesOptions.Batch);
        }

        /// <summary>
        /// Saves the changes.
        /// </summary>
        /// <param name="options">Flags for controlling the behavior of the save.</param>
        public void SaveChanges(SaveChangesOptions options)
        {
            IEnumerable<string> tables = this.Entities.Select(e => this.GetTableInfo(e.Entity.GetType()).TableName).Distinct();

            foreach (string table in tables)
            {
                this.DoEntityRequest(options, table);
            }

            if (this.RemovedEntities.Count > 0)
            {
                this.RemovedEntities.ForEach(re => this.Entities.Remove(re));
                this.RemovedEntities.Clear();
            }
        }

        /// <summary>
        /// Creates the TTableStorageEngine engine.
        /// </summary>
        /// <param name="tableName">Name of the table.</param>
        /// <returns>A new TTableStorageEngine for the specified table</returns>
        public TTableStorageEngine CreateEngine(string tableName)
        {
            return new TTableStorageEngine(
                this.ServiceRoot,
                this.Credentials,
                fTableNameTableInfo[tableName],
                tableName,
                this.RetryPolicy,
                this.IncrementTransactionCount);
        }

        /// <summary>
        /// Creates a query on the table associated with <typeparamref name="EntityT"/>
        /// </summary>
        /// <typeparam name="EntityT">The type of entity contained in the table.</typeparam>
        /// <returns>An IQueryable against the table associated with <typeparamref name="EntityT"/></returns>
        public TTableStorageQuery<EntityT> CreateQuery<EntityT>()
        {
            TTableStorageQuery<EntityT> qry = new TTableStorageQuery<EntityT>(
                this.GetTableInfo(typeof(EntityT)),
                tableName => this.CreateEngine(tableName),
                t =>
                {
                    TTablePropertyInfo tablePropInfo = this.GetTableInfo(typeof(EntityT)).PropertyList.Where(tpi => tpi.Aliases.Contains(t.Name)).FirstOrDefault();

                    if (tablePropInfo == null)
                    {
                        throw new TTableStorageException(String.Format("Property '{0}' not found in table info for type {1}.  This may be due to using a non-persisted wrapper around a persisted property.",
                            t.Name, typeof(EntityT)));
                    }

                    return tablePropInfo.FieldName;
                }, this);

            return qry;
        }

        /// <summary>
        /// Gets the name of the table.
        /// </summary>
        /// <param name="entityType">Type of the entity.</param>
        /// <returns>The name of the table</returns>
        private TTableStorageInfo GetTableInfo(Type entityType)
        {
            TTableStorageInfo typeTableInfoResult;

            if (!fTypeTableInfo.TryGetValue(entityType, out typeTableInfoResult))
            {
                typeTableInfoResult = this.DoGetTableInfo(entityType);

                if (entityType.IsAbstract)
                {
                    //// TODO: add support for KnownTypeAttribute lookup method.

                    foreach (Type knownType in entityType.GetCustomAttributes(typeof(KnownTypeAttribute), false).Where(attr => ((KnownTypeAttribute)attr).Type != null).Select(attr => ((KnownTypeAttribute)attr).Type))
                    {
                        if (knownType == entityType)
                        {
                            throw new TTableStorageException("It appears you have listed an abstract type as it's own KnownType.  KnownTypes should be concrete implementations of abstract types.");
                        }

                        this.GetTableInfo(knownType);
                    }
                }

                lock (fTableInfoLockObject)
                {
                    // someone else may have added it to the Dictionary while
                    // we were looking it up ... so we need to recheck
                    TTableStorageInfo currentTableInfo = null;

                    if (!fTypeTableInfo.TryGetValue(entityType, out currentTableInfo))
                    {
                        fTypeTableInfo.Add(entityType, typeTableInfoResult);
                    }

                    List<TTableStorageInfo> tableNameTypeListResult;
                    if (!fTableNameTableInfo.TryGetValue(typeTableInfoResult.TableName, out tableNameTypeListResult))
                    {
                        tableNameTypeListResult = new List<TTableStorageInfo>() { typeTableInfoResult };
                        fTableNameTableInfo.Add(typeTableInfoResult.TableName, tableNameTypeListResult);
                    }
                    else
                    {
                        if (!tableNameTypeListResult.Contains(typeTableInfoResult))
                        {
                            tableNameTypeListResult.Add(typeTableInfoResult);
                        }
                    }
                }

                if (typeof(ITableEntity).IsAssignableFrom(entityType.BaseType) && entityType.BaseType.IsAbstract && !fTypeTableInfo.ContainsKey(entityType.BaseType))
                {
                    this.GetTableInfo(entityType.BaseType);
                }
            }

            return typeTableInfoResult;
        }

        /// <summary>
        /// Adds a new entity if necessary and changes the state to <paramref name="state"/>.
        /// </summary>
        /// <param name="entity">The entity to add.</param>
        /// <param name="state">The state of the entity.</param>
        /// <returns>The EntityInfo associated with the entity.</returns>
        private TTableEntityTransactionInfo AddOrChangeState(ITableEntity entity, EntityStates state)
        {
            TTableEntityTransactionInfo info;

            if (!this.IsTrackingEntity(entity))
            {
                info = new TTableEntityTransactionInfo(entity, state);

                this.Entities.Add(info);
            }
            else
            {
                info = this.Entities.SingleOrDefault(ei => ei.Entity == entity);

                if (info.State == EntityStates.Added && state == EntityStates.Deleted)
                {
                    this.Entities.Remove(info);
                    return null;
                }

                info.State = state;
            }

            switch (info.State)
            {
                case EntityStates.Added:
                    info.Verb = HttpConstants.Post;
                    info.SuccessCode = HttpStatusCode.Created;
                    info.OnSuccess = this.OnSuccess_AddOrUpdate;
                    info.OnFailure = this.OnFailure;
                    break;
                case EntityStates.Modified:
                    info.Verb = HttpConstants.Put;
                    info.SuccessCode = HttpStatusCode.NoContent;
                    info.OnSuccess = this.OnSuccess_AddOrUpdate;
                    info.OnFailure = this.OnFailure;
                    break;
                case EntityStates.Deleted:
                    info.Verb = HttpConstants.Delete;
                    info.SuccessCode = HttpStatusCode.NoContent;
                    info.OnSuccess = this.OnSuccess_Delete;
                    info.OnFailure = this.OnFailure;
                    break;
            }

            return info;
        }

        /// <summary>
        /// Called when a TTableEntityTransactionInfo is successfully added or updated
        /// </summary>
        /// <param name="info">The TTableEntityTransactionInfo containing the entity that was successfully updated and details of the transaction.</param>
        private void OnSuccess_AddOrUpdate(TTableEntityTransactionInfo info)
        {
            info.State = EntityStates.Unchanged;
        }

        /// <summary>
        /// Called when a TTableEntityTransactionInfo is successfully deleted
        /// </summary>
        /// <param name="info">The TTableEntityTransactionInfo containing the entity that was successfully deleted and details of the transaction.</param>
        private void OnSuccess_Delete(TTableEntityTransactionInfo info)
        {
            this.RemovedEntities.Add(info);
        }

        /// <summary>
        /// Called when a TTableEntityTransactionInfo fails to perform an add/update/delete
        /// </summary>
        /// <param name="info">The TTableEntityTransactionInfo containing the entity that failed to be deleted and details of the transaction.</param>
        private void OnFailure(TTableEntityTransactionInfo info)
        {
            ////throw info.Error;
        }

        /// <summary>
        /// Performs a Table Service request on entities being tracked by the context.
        /// </summary>
        /// <param name="options">The save options (None, Batch, etc.) <see cref="SaveChangesOptions"/></param>
        /// <param name="tableName">Name of the table used in the request</param>
        private void DoEntityRequest(SaveChangesOptions options, string tableName)
        {
            // ToList is used to create a copy of targets, because processing the results of the requests will probably change
            // the resulting content of targets.
            IEnumerable<TTableEntityTransactionInfo> targets = this.Entities
                .Where(
                    targetInfo =>
                    {
                        TTableStorageInfo tableInfo = this.GetTableInfo(targetInfo.Entity.GetType());
                        return targetInfo.State != EntityStates.Unchanged && tableInfo.TableName == tableName;
                    }).ToList();

            HttpWebRequest request;

            List<TTableStorageInfo> tableInfosForTargets = fTableNameTableInfo[tableName];

            TTableStorageEngine engine = new TTableStorageEngine(
                this.ServiceRoot,
                this.Credentials,
                tableInfosForTargets,
                tableName,
                this.RetryPolicy,
                this.IncrementTransactionCount);

            if (options.HasFlag(SaveChangesOptions.Batch))
            {
                int targetCount = targets.Count();

                if (targetCount != 0)
                {
                    foreach (IEnumerable<TTableEntityTransactionInfo> splitTargets in targets.Split(TableStorageConstants.TransactionLimit))
                    {
                        Retry.Invoke(
                            () =>
                            {
                                request = engine.CreateEntityBatchRequest(splitTargets);
                                engine.ReadEntityBatchResponse(request, splitTargets);
                            }, this.RetryPolicy);
                    }
                }
            }
            else
            {
                foreach (TTableEntityTransactionInfo target in targets)
                {
                    List<TTableEntityTransactionInfo> truncTargets = new List<TTableEntityTransactionInfo>() { target };

                    Retry.Invoke(
                        () =>
                        {
                            request = engine.CreateEntityBatchRequest(truncTargets);
                            engine.ReadEntityBatchResponse(request, truncTargets);
                        }, this.RetryPolicy);
                }
            }
        }

        /// <summary>
        /// Increments the transaction count.
        /// </summary>
        private void IncrementTransactionCount()
        {
            ++this.fTransactionCount;
            
            if (this.IncrementingTransactionCount != null)
            {
                this.IncrementingTransactionCount();
            }
        }
    }
}
